// standard headers
#include <istream>
#include <vector>
#include <string>
#include <fstream>
#include <sstream>
#include <iostream>
#include <algorithm>
#include <stdexcept>
#include <map>
#include <cstdlib>

// boost headers
#include <boost/algorithm/string.hpp>
#include <boost/foreach.hpp>


// components headers
#include <tree.h>
#include <matrix_io.h>
#include <tree_ops.h>
#include <portable_timer.h>
#include <index_string.h>
#include <progress_indicator.h>
#include <print_error.h>
#include <ultoa.h>
#define DEBUG_LEVEL 6
#include <debug_func.h>


#include "get_arg.h"



using std::vector;
using std::ostringstream;

using std::string;
using std::fill;
using std::ofstream;
using std::cout;
using std::cin;
using std::cerr;
using std::ifstream;
using std::map;
using std::pair;
using std::istream;
using std::max;
using std::min;
using namespace bcpp_matrix;
using namespace bcpp_tree_ops;


struct do_reroot
{
	vector<bcpp_tree::t_edge>	branches;
	std::ostream&				os;
	bcpp_tree::t_tree&			tree;
	op_print_tree_elaborate&	print_tree_elaborate;
	do_reroot(	bcpp_tree::t_tree&				tree_,
				std::ostream& 						os_,
				op_print_tree_elaborate&			print_tree_elaborate_):
			os(os_),
			tree(tree_),
			print_tree_elaborate(print_tree_elaborate_){}
	void recurse(bcpp_tree::t_node* node, unsigned);
};

void do_reroot::recurse(bcpp_tree::t_node* node, unsigned ignore)
{                                                                      ;
	debug_func_cerr(5);
	unsigned cnt_children = node->get_children_count();
	// is leaf
	if (!cnt_children)
	{
		return;
	}

	// handle all  child with grandchildren
	for (unsigned i = 0; i < cnt_children; ++i)
	{
		if (i == ignore)
			continue;
		if (node->get_child(i)->get_children_count())
		{
			tree.reroot(i, 0);
			branches.push_back(bcpp_tree::t_edge(tree.get_root()->get_child(0), tree.get_root()->get_child(1)));
			print_tree_elaborate.print();
			recurse(node->get_child(i),0);
			tree.undo_reroot(i, 0);

			tree.reroot(i, 1);
			branches.push_back(bcpp_tree::t_edge(tree.get_root()->get_child(0), tree.get_root()->get_child(1)));
			print_tree_elaborate.print();
			recurse(node->get_child(i),1);
			tree.undo_reroot(i, 1);
		}
	}
}

struct do_evaluate_tree
{
	bcpp_tree::t_tree&			tree;

	// mapping tables
	t_matrix_io_template<unsigned>& duplications_map;
	t_matrix_io_template<unsigned>& deletions_map;
	t_matrix_io_template<unsigned>& transitions_map;
	t_matrix_io_template<unsigned>& speciation_nodes_map;
	t_index_string&				gene_names;
	vector<unsigned>&			genes_to_species;
	unsigned					unknown_species_index;

	// state
	std::vector<unsigned>		duplicates;
	std::vector<unsigned>		deletions;
	std::vector<unsigned>		cnt_duplicates;
	std::vector<unsigned>		cnt_deletions;
	std::vector<unsigned>		node_type;
	std::vector<unsigned>		speciation;

	// constructor
	do_evaluate_tree(	bcpp_tree::t_tree&	tree_,
						t_matrix_io_template<unsigned>& duplications_map_,
						t_matrix_io_template<unsigned>& deletions_map_,
						t_matrix_io_template<unsigned>& transitions_map_,
						t_matrix_io_template<unsigned>& speciation_nodes_map_,
						t_index_string&					gene_names_,
						vector<unsigned>&				genes_to_species_
						)
					:
						tree(tree_),
						duplications_map	(duplications_map_		),  
						deletions_map		(deletions_map_			), 
						transitions_map		(transitions_map_ 		), 
						speciation_nodes_map(speciation_nodes_map_	), 
						gene_names			(gene_names_			),
						genes_to_species	(genes_to_species_		),
						unknown_species_index(duplications_map.size() -1),
		
						duplicates			(tree.get_node_count(),0),
						deletions 			(tree.get_node_count(),0),
						cnt_duplicates		(tree.get_node_count(),0),
						cnt_deletions 		(tree.get_node_count(),0),
						node_type  			(tree.get_node_count(),0),
						speciation			(tree.get_node_count(),0)
	{
	}
	void first_pass()
	{
		unsigned root_UID = recurse(tree.get_root());
		cout << "duplications=\t" << cnt_duplicates[root_UID] << "\n";
		cout << "deletions=\t" << cnt_deletions[root_UID] << "\n";
	}
	unsigned first_pass_recurse(bcpp_tree::t_node* node);
	void reroot_recurse(bcpp_tree::t_node* node, unsigned ignore);
};

void do_evaluate_tree::reroot_recurse(bcpp_tree::t_node* node, unsigned ignore)
{
	debug_func_cerr(5);
	unsigned cnt_children = node->get_children_count();

	// is leaf: cannot root here
	if (!cnt_children)
		return;

	// handle all  child with grandchildren
	for (unsigned i = 0; i < cnt_children; ++i)
	{
		if (i == ignore)
			continue;
		if (node->get_child(i)->get_children_count())
		{
			tree.reroot(i, 0);
			branches.push_back(bcpp_tree::t_edge(tree.get_root()->get_child(0), tree.get_root()->get_child(1)));
			reroot_recurse(node->get_child(i),0);

			tree.undo_reroot(i, 0);

			tree.reroot(i, 1);
			branches.push_back(bcpp_tree::t_edge(tree.get_root()->get_child(0), tree.get_root()->get_child(1)));
			reroot_recurse(node->get_child(i),1);
			tree.undo_reroot(i, 1);
		}
	}
}

unsigned do_evaluate_tree::first_pass_recurse(bcpp_tree::t_node* node)
{                                                                      ;
	debug_func_cerr(5);
	unsigned UID = node->get_UID();
	assert(UID < tree.get_node_count());
	unsigned cnt_children = node->get_children_count();
	//
	// is leaf
	// 	look up species of leaf, otherwise unknown
	if (!cnt_children)
	{
		unsigned gene_index;
		if(!gene_names.is_indexed(node->get_name(), gene_index))
			node_type[UID] = unknown_species_index;
		else
		{
			cerr << genes_to_species.size() << "\n";
			assert(gene_index < genes_to_species.size());
			node_type[UID] = genes_to_species[gene_index];
		}
		return UID;
	}

	// handle all  child with grandchildren
	unsigned UID1 = recurse(node->get_child(0));
	unsigned UID2 = recurse(node->get_child(1));

	unsigned nt1 =  node_type[UID1];
	unsigned nt2 =  node_type[UID2];
	node_type[UID] = transitions_map.matrix[nt1][nt2];
	duplicates[UID] = duplications_map.matrix[nt1][nt2];
	cnt_duplicates[UID] = cnt_duplicates[UID1] + cnt_duplicates[UID2] + duplicates[UID];
	deletions[UID] = deletions_map.matrix[nt1][nt2];
	cnt_deletions[UID] = cnt_deletions[UID1] + cnt_deletions[UID2] + deletions[UID];
	return UID;
}







//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

//	main

//8888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

int main (int argc, char *argv[])
{
	// we use cout to print progress so flush per call
//	cout.setf(std::ios_base::unitbuf);

	std::ios::sync_with_stdio(false);

	t_program_args args;
	try
	{
		if (!process_command_line_options(argc, argv, args))
			return 1;


		//
		//	Reading species tables
		//
		args.VERBOSE() << "Read species tables...\n";
		t_matrix_io_template<unsigned> duplications;
		t_matrix_io_template<unsigned> deletions;
		t_matrix_io_template<unsigned> transitions;
		t_matrix_io_template<unsigned> speciation_nodes;
		t_index_string species;
		{
			ifstream ifs(args.species_tables.c_str());
			if (!ifs)
				throw std::runtime_error("Could not open " + args.species_tables);
	
			unsigned ignore_cnt_duplicate_pairs;
			read_matrix(ifs, duplications    ,  eMIO_SQUARE, ignore_cnt_duplicate_pairs);
			read_matrix(ifs, deletions       ,  eMIO_SQUARE, ignore_cnt_duplicate_pairs);
			read_matrix(ifs, transitions     ,  eMIO_SQUARE, ignore_cnt_duplicate_pairs);
			read_matrix(ifs, speciation_nodes,  eMIO_SQUARE, ignore_cnt_duplicate_pairs);
			BOOST_FOREACH(const string& s, duplications.names) 
				species.index(s);
		}
		
		//
		//	Reading genes to species map
		//
		args.VERBOSE() << "Reading genes to species map...\n";
		t_index_string gene_names;
		vector<unsigned> genes_to_species;
		{
			ifstream ifs(args.gene_to_species.c_str());
			if (!ifs)
				throw std::runtime_error("Could not open " + args.gene_to_species);
			unsigned cnt_lines = 0;
			while (ifs)
			{
				string line;
				getline(ifs, line);
				++cnt_lines;

				// ignore blank lines
				if (!line.length())
					continue;

				//vector<boost::iterator_range<std::string::iterator> >
				vector<string> fields;
				boost::split(fields, line, std::bind2nd(std::equal_to<char>(), '\t'));

				if (fields.size() != 2)
					throw  std::runtime_error("Line not tab-delimited on #" + ultoa(cnt_lines));


				// set strange species to unknown
				unsigned species_index;
				if(!species.is_indexed(fields[1], species_index))
					species_index = species.size() - 1;

				bool gene_added;
				gene_names.index(fields[0], gene_added);
				if (gene_added)
					genes_to_species.push_back(species_index);
			}
		}

		
		
		//write_matrix(cout, duplications    ,  eMIO_SQUARE, false);
		//write_matrix(cout, deletions       ,  eMIO_SQUARE, false);
		//write_matrix(cout, transitions     ,  eMIO_SQUARE, false);
		//write_matrix(cout, speciation_nodes,  eMIO_SQUARE, false);

		// read gene tree
		bcpp_tree::t_tree initial_tree;
		unsigned line_cnt;
		read_tree_from_stream(cin, initial_tree, line_cnt);
		if (!cin)
			throw std::runtime_error("No tree supplied.");


		do_evaluate_tree evaluate_tree(	initial_tree,
							duplications,
							deletions,
							transitions,
							speciation_nodes,
							gene_names,
							genes_to_species);
		evaluate_tree.evaluate();
		
	
/*		cerr << "\tCalculating branch lengths...\n";
		// calculate branch lengths
		bcpp_tree_ops::do_calculate_branch_lengths branch_lengths(initial_tree);
	
		// set node type to none
		std::vector<eNODETYPE>	node_types(initial_tree.get_node_count(), eNONE);
	
	
		// print
		op_print_tree_elaborate print_tree_elaborate
									(initial_tree,
									cout,
									branch_lengths.lengths_to_leaves,
									branch_lengths.lengths_to_root,
									node_types,
									3	// precision
									);
		cout << "Original tree:\n";
		print_tree_elaborate.print();
		do_reroot reroot(initial_tree, cout, print_tree_elaborate);
		cout << "Rerooting recursively...\n";
		reroot.recurse(initial_tree.get_root(),0);
		reroot.recurse(initial_tree.get_root(),1);
		cerr << "Rerooting randomly...\n";
		cout << "Rerooting randomly...\n";
		for (unsigned i = 0; i < 150; ++i)
		{
			unsigned index = rand() % reroot.branches.size();
			initial_tree.reroot(reroot.branches[index].parent, reroot.branches[index].child);
			print_tree_elaborate.print();
			initial_tree.assert_consistency();
		}
	
*/	
		cerr << "\n\tAll finished!!\n\n";
	}
	catch (std::exception& e)
	{
		std_print_error(cerr,  e.what());
		return 1;
	}

	return 0;
}


